package com.starhub.application.rag.processor;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;

import java.util.List;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;

import com.starhub.application.collection.entity.KnowledgeCollection;
import com.starhub.application.interfaces.model.ModelConfigService;
import com.starhub.application.model.enums.VectorModelTypeEnum;
import com.starhub.application.model.util.DimensionUtil;
import com.starhub.application.rag.constant.RagConstant;
import com.starhub.application.rag.dto.MetaRagDto;
import com.starhub.application.rag.dto.RetrieveDto;
import com.starhub.application.rag.result.RagResult;
import com.starhub.application.rag.store.RetrieveStore;
import com.starhub.common.bean.model.ModelConfig;

@Slf4j
public class EmbeddingProcessor {

    @Autowired
    private RetrieveStore retrieveStore;

    @Autowired
    private ModelConfigService modelConfigService;

    /**
     * 执行单个集合的检索任务
     * 
     * @param collection        知识库集合
     * @param metaRagDto        元数据检索请求
     * @param vectorModel       向量模型标识
     * @param questionEmbedding 问题向量
     * @return 检索结果
     */
    protected List<String> embeddingRetrieval(KnowledgeCollection collection,
            MetaRagDto metaRagDto,
            String vectorModel,
            Embedding questionEmbedding) {
        try {
            // 构建检索请求 相当于把用户输入在当前元数据表中进行向量检索
            RetrieveDto retrieveDto = RetrieveDto.builder()
                    // .queryText(metaRagDto.getQueryText())
                    .vectorModel(vectorModel)
                    .collectionIdentifier(collection.getCollectionIdentifier())
                    .dimension(DimensionUtil.getDimension(vectorModel))
                    .topK(RagConstant.RETRIEVE_EMBEDDING_COUNT)
                    .minScore(RagConstant.RETRIEVE_EMBEDDING_MIN_SCORE)
                    .build();

            // 执行检索
            RagResult ragResult = retrieveStore.retrieve(retrieveDto, questionEmbedding);

            // 将检索结果转换为知识条目
            List<String> itemIds = ragResult.getSearchResult().matches().stream()
                    .map(match -> {
                        TextSegment textSegment = match.embedded();
                        Metadata metadata = textSegment.metadata();
                        return metadata.getString(RagConstant.EmbeddingMetaData.ITEM_ID);
                    })
                    .collect(Collectors.toList());

            return itemIds;
        } catch (Exception e) {
            log.error("集合 {} 检索失败: {}", collection.getCollectionIdentifier(), e.getMessage(), e);
            return null;
        }
    }
}
